import os
import sys
import random
import argparse
import numpy as np
from PIL import Image, ImageFile



BLUE_IMAGE_PATH = r"C:\Users\kiwi\Desktop\face-mask-detector\Data_Generator\images\blue-mask.png"

def create_mask(image_path):
    pic_path = image_path
    
    mask_path = BLUE_IMAGE_PATH
    # mask_path = "/media/preeth/Data/prajna_files/mask_creator/face_mask/images/blue-mask.png"
    show = False
    model = "hog"
    FaceMasker(pic_path, mask_path, show, model).mask()  #mask_path口罩的路径



class FaceMasker:
    KEY_FACIAL_FEATURES = ('nose_bridge', 'chin')

    def __init__(self, face_path, mask_path, show=False, model='hog'):
        self.face_path = face_path  #图片模型
        self.mask_path = mask_path  #口罩模型
        self.show = show
        self.model = model
    #     self._face_img: ImageFile = None
    #     self._mask_img: ImageFile = None

    def mask(self):
        import face_recognition

        face_image_np = face_recognition.load_image_file(self.face_path) #加载图片
        face_locations = face_recognition.face_locations(face_image_np, model=self.model)  #查找图片中人脸（上下左右）的位置，图像中可能有多个人脸
        face_landmarks = face_recognition.face_landmarks(face_image_np, face_locations)  #返回人的64个特征点的坐标
        #实现array到image的转换,返回image
        self._face_img = Image.fromarray(face_image_np)  #将图片转化为np.array形式
        self._mask_img = Image.open(self.mask_path)      #打开mask图片

        found_face = False
        for face_landmark in face_landmarks:  #针对每一张图片
            # check whether facial features meet requirement
            skip = False
            for facial_feature in self.KEY_FACIAL_FEATURES: #判断每一张图片中的每一个特征点，是否为鼻子和脸颊
                if facial_feature not in face_landmark:
                    skip = True
                    break
            if skip:
                continue

            # mask face
            found_face = True
            self._mask_face(face_landmark)

        if found_face:
            if self.show:
                self._face_img.show()

            # save
            self._save()
        else:
            print('Found no face.')

    def _mask_face(self, face_landmark: dict):
        nose_bridge = face_landmark['nose_bridge']  #鼻子
        nose_point = nose_bridge[len(nose_bridge) * 1 // 4]  #鼻子的最下边的点
        nose_v = np.array(nose_point)

        chin = face_landmark['chin']  #脸颊
        chin_len = len(chin)
        chin_bottom_point = chin[8]  #中间最下面的点
        # chin_bottom_point = chin[chin_len // 2]
        chin_bottom_v = np.array(chin_bottom_point)
        # chin_left_point = chin[chin_len // 8]
        # chin_right_point = chin[chin_len * 7 // 8]
        chin_left_point = chin[2]  #最左边的点
        chin_right_point = chin[16]  #最右边的点

        # split mask and resize
        width = self._mask_img.width  #口罩原始宽度
        height = self._mask_img.height #口罩原始长度
        width_ratio = 1.2
        new_height = int(np.linalg.norm(nose_v - chin_bottom_v))

        # 自适应对左脸进行适配
        mask_left_img = self._mask_img.crop((0, 0, width // 2, height))  #左半边口罩
        mask_left_width = self.get_distance_from_point_to_line(chin_left_point, nose_point, chin_bottom_point)  #宽的距离
        mask_left_width = int(mask_left_width * width_ratio)  #乘以比率获得将半边脸全部遮住
        mask_left_img = mask_left_img.resize((mask_left_width, new_height))  #将左半边的口罩进行ersize，以适应半边脸

        # 相同的操作对右脸进行适配
        mask_right_img = self._mask_img.crop((width // 2, 0, width, height))
        mask_right_width = self.get_distance_from_point_to_line(chin_right_point, nose_point, chin_bottom_point)
        mask_right_width = int(mask_right_width * width_ratio)
        mask_right_img = mask_right_img.resize((mask_right_width, new_height))

        # merge mask
        size = (mask_left_img.width + mask_right_img.width, new_height)
        mask_img = Image.new('RGBA', size) #创建一幅给定模式（mode）和尺寸（size）的图片
        mask_img.paste(mask_left_img, (0, 0), mask_left_img)  #图片指定左半边区域替换
        mask_img.paste(mask_right_img, (mask_left_img.width, 0), mask_right_img)  ##图片指定右半边区域替换

        # rotate mask
        angle = np.arctan2(chin_bottom_point[1] - nose_point[1], chin_bottom_point[0] - nose_point[0])
        rotated_mask_img = mask_img.rotate(angle, expand=True)

        # calculate mask location
        center_x = (nose_point[0] + chin_bottom_point[0]) // 2
        center_y = (nose_point[1] + chin_bottom_point[1]) // 2

        offset = mask_img.width // 2 - mask_left_img.width
        radian = angle * np.pi / 180
        box_x = center_x + int(offset * np.cos(radian)) - rotated_mask_img.width // 2
        box_y = center_y + int(offset * np.sin(radian)) - rotated_mask_img.height // 2

        # add mask
        self._face_img.paste(mask_img, (box_x, box_y), mask_img)

    def _save(self):
    	# os.path.splitext(“文件路径”)    分离文件名与扩展名；默认返回(fname,fextension)元组，可做分片操作
        path_splits = os.path.splitext(self.face_path)
        new_face_path = path_splits[0] + '-with-mask' + path_splits[1]
        self._face_img.save(new_face_path)
        print(f'Save to {new_face_path}')

    @staticmethod
    def get_distance_from_point_to_line(point, line_point1, line_point2):   #三个点之间距离的计算，得到宽的大小
        distance = np.abs((line_point2[1] - line_point1[1]) * point[0] +
                          (line_point1[0] - line_point2[0]) * point[1] +
                          (line_point2[0] - line_point1[0]) * line_point1[1] +
                          (line_point1[1] - line_point2[1]) * line_point1[0]) / \
                   np.sqrt((line_point2[1] - line_point1[1]) * (line_point2[1] - line_point1[1]) +
                           (line_point1[0] - line_point2[0]) * (line_point1[0] - line_point2[0]))
        return int(distance)


if __name__ == '__main__':
    #cli()
    create_mask(image_path)
